/**
 * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2019. ALL RIGHTS RESERVED.
 *
 * See file LICENSE for terms.
 */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <ucm/cncl/cnclmem.h>

#include <ucm/event/event.h>
#include <ucm/mmap/mmap.h>
#include <ucm/util/log.h>
#include <ucm/util/reloc.h>
#include <ucm/util/replace.h>
#include <ucm/util/sys.h>
#include <ucs/sys/compiler.h>
#include <ucs/sys/preprocessor.h>

#include <sys/mman.h>
#include <string.h>

// cnrtErrorMax
// cnrtRet_t cnrtMalloc(void **pPtr, size_t bytes);
UCM_DEFINE_REPLACE_DLSYM_FUNC(cnrtMalloc, cnrtRet_t,
                              cnrtErrorMax, void**, size_t)
/*
ucm_orig_cnrtMalloc         origin dlsym function
ucm_override_cnrtMalloc     calls ucm_cnrtMalloc
*/
// cnrtRet_t cnrtHostMalloc(void **pPtr, size_t bytes);
UCM_DEFINE_REPLACE_DLSYM_FUNC(cnrtHostMalloc, cnrtRet_t,
                              cnrtErrorMax, void**, size_t)
// cnrtRet_t cnrtFree(void *ptr);
UCM_DEFINE_REPLACE_DLSYM_FUNC(cnrtFree, cnrtRet_t,
                              cnrtErrorMax, void*)
// cnrtRet_t cnrtFreeHost(void *ptr);
UCM_DEFINE_REPLACE_DLSYM_FUNC(cnrtFreeHost, cnrtRet_t,
                              cnrtErrorMax, void*)

static UCS_F_ALWAYS_INLINE void
ucm_dispatch_mem_type_alloc(void *addr, size_t length, ucs_memory_type_t mem_type)
{
    ucm_event_t event;

    event.mem_type.address  = addr;
    event.mem_type.size     = length;
    event.mem_type.mem_type = mem_type;
    ucm_event_dispatch(UCM_EVENT_MEM_TYPE_ALLOC, &event);
}

static UCS_F_ALWAYS_INLINE void
ucm_dispatch_mem_type_free(void *addr, size_t length, ucs_memory_type_t mem_type)
{
    ucm_event_t event;
    cnrtRet_t status;
    cnrtPointerAttributes_t attr;

    if (addr == NULL) {
        return;
    }

    if (length==0) {
        status = cnrtPointerGetAttributes(&attr, addr);
        if (status != cnrtSuccess) {
            ucm_trace("cnrtPointerGetAttributes failed.");
            return;
        }
        length = attr.size;

    }

    event.mem_type.address  = addr;
    event.mem_type.size     = length;
    event.mem_type.mem_type = mem_type;
    ucm_event_dispatch(UCM_EVENT_MEM_TYPE_FREE, &event);
}

cnrtRet_t ucm_cnrtMalloc(void **pPtr, size_t bytes)
{
    cnrtRet_t status;

    ucm_event_enter();

    status = ucm_orig_cnrtMalloc(pPtr, bytes);
    if (status != cnrtSuccess) {
        ucm_trace("ucm_cnrtMalloc(ptr=%p size:%lu)", *pPtr, bytes);
        ucm_dispatch_mem_type_alloc(*pPtr, bytes, UCS_MEMORY_TYPE_UNKNOWN);
    }

    ucm_event_leave();
    return status;
}

cnrtRet_t cnrtHostMalloc(void **pPtr, size_t bytes)
{
    cnrtRet_t status;

    ucm_event_enter();

    status = ucm_orig_cnrtHostMalloc(pPtr, bytes);
    if (status != cnrtSuccess) {
        ucm_trace("ucm_cnrtHostMalloc(ptr=%p size:%lu)", *pPtr, bytes);
        ucm_dispatch_mem_type_alloc(*pPtr, bytes, UCS_MEMORY_TYPE_UNKNOWN);
    }

    ucm_event_leave();
    return status;
}

cnrtRet_t ucm_cnrtFree(void *ptr)
{
    cnrtRet_t status;

    ucm_event_enter();

    ucm_trace("ucm_cnrtFree(ptr=%p)", ptr);

    ucm_dispatch_mem_type_free(ptr, 0, UCS_MEMORY_TYPE_CNCL);

    status = ucm_orig_cnrtFree(ptr);

    ucm_event_leave();
    return status;
}

cnrtRet_t ucm_cnrtFreeHost(void *ptr)
{
    cnrtRet_t status;

    ucm_event_enter();

    ucm_trace("ucm_cnrtFreeHost(ptr=%p)", ptr);

    ucm_dispatch_mem_type_free(ptr, 0, UCS_MEMORY_TYPE_HOST);

    status = ucm_orig_cnrtFreeHost(ptr);

    ucm_event_leave();
    return status;
}


static ucm_reloc_patch_t cnclpatches[] = {
    {UCS_PP_MAKE_STRING(cnrtMalloc),    // page-locked memory
     ucm_override_cnrtMalloc},
    {UCS_PP_MAKE_STRING(cnrtHostMalloc),     // page-locked memory
     ucm_override_cnrtHostMalloc},
    {UCS_PP_MAKE_STRING(cnrtFree),
     ucm_override_cnrtFree},
    {UCS_PP_MAKE_STRING(cnrtFreeHost),
     ucm_override_cnrtFreeHost},
    {NULL, NULL}
};

static ucs_status_t ucm_cnclmem_install(int events)
{
    static pthread_mutex_t install_mutex = PTHREAD_MUTEX_INITIALIZER;
    static int ucm_cnclmem_installed = 0;
    ucm_reloc_patch_t *patch;
    ucs_status_t status                  = UCS_OK;

    if (!(events & (UCM_EVENT_MEM_TYPE_ALLOC | UCM_EVENT_MEM_TYPE_FREE))) {
        goto out;
    }

    pthread_mutex_lock(&install_mutex);

    if (ucm_cnclmem_installed) {
        goto out_unlock;
    }

    for (patch = cnclpatches; patch->symbol != NULL; ++patch) {
        status = ucm_reloc_modify(patch);
        if (status != UCS_OK) {
            ucm_warn("failed to install relocation table entry for '%s'", patch->symbol);
            goto out_unlock;
        }
    }

    ucm_info("cncl hooks are ready");
    ucm_cnclmem_installed = 1;

out_unlock:
    pthread_mutex_unlock(&install_mutex);
out:
    return status;
}

static int ucm_cncl_scan_regions_cb(void *arg, void *addr, size_t length,
                                       int prot, const char *path)
{
    static const char *cncl_path_pattern = "/dev/cambricon";
    ucm_event_handler_t *handler         = arg;
    ucm_event_t event;

    if ((prot & (PROT_READ | PROT_WRITE | PROT_EXEC)) &&
        strncmp(path, cncl_path_pattern, strlen(cncl_path_pattern))) {
        return 0;
    }

    ucm_debug("dispatching initial memtype allocation for %p..%p %s", addr,
              UCS_PTR_BYTE_OFFSET(addr, length), path);

    event.mem_type.address  = addr;
    event.mem_type.size     = length;
    event.mem_type.mem_type = UCS_MEMORY_TYPE_LAST; /* unknown memory type */

    ucm_event_enter();
    handler->cb(UCM_EVENT_MEM_TYPE_ALLOC, &event, handler->arg);
    ucm_event_leave();

    return 0;
}

static void ucm_cnclmem_get_existing_alloc(ucm_event_handler_t *handler)
{
    if (handler->events & UCM_EVENT_MEM_TYPE_ALLOC) {
        ucm_parse_proc_self_maps(ucm_cncl_scan_regions_cb, handler);
    }
}

static ucm_event_installer_t ucm_cncl_initializer = {
    .install            = ucm_cnclmem_install,
    .get_existing_alloc = ucm_cnclmem_get_existing_alloc
};

UCS_STATIC_INIT
{
    ucs_list_add_tail(&ucm_event_installer_list, &ucm_cncl_initializer.list);
}

UCS_STATIC_CLEANUP
{
    ucs_list_del(&ucm_cncl_initializer.list);
}
